#! /usr/bin/env python
# coding=utf-8


# 读取simulation的结果


pep_cut_ali = r'pep_cut_ali.fa'

seq_file_name_aa = 'pep_cut_ali.aa'

tree_file_name = 'spe_tree'
aactl_name = 'aaml.ctl'


ncatG = 4


from Bio import SeqIO
import os
import shutil
import re
import argparse
import sys




parser = argparse.ArgumentParser(
    description='''提取simulation文件 为paml 做准备
    用法:
    conver04_read_simulation.py -i OGS_for_second_step -n 4 -d jones.dat -o simulate
    由大天才于2021年11月11日创建于浙江农业大学''')





parser.add_argument('-i',
                help='必须给定，OGS的工作目录')

parser.add_argument('-n',
                help='ncatG,必须给定')

parser.add_argument('-d',
                help='aaml需要的替换模型文件，必须给定')

parser.add_argument('-o',
                help='输出用于跑paml的路径，必须给定')

args = parser.parse_args()

if not args.i or not args.n or not args.d or not args.o:
    parser.print_help()
    sys.exit()



infile = args.i

ncatG = int(args.n)


dat_file_path  = args.d

#infile = 'test'
outfile = args.o


try:
	shutil.rmtree(outfile)
except:
	pass
try:
	os.mkdir(outfile)
except:
	pass




paml_tmp = '''

      seqfile = %s
     treefile = %s

      outfile = mlc           * main result file name
        noisy = 3  * 0,1,2,3,9: how much rubbish on the screen
      verbose = 1  * 0: concise; 1: detailed, 2: too much
      runmode = 0  * 0: user tree;  1: semi-automatic;  2: automatic
                   * 3: StepwiseAddition; (4,5):PerturbationNNI; -2: pairwise

      seqtype = 2  * 1:codons; 2:AAs; 3:codons-->AAs
   aaRatefile = my.dat * only used for aa seqs with model=empirical(_F)
                   * dayhoff.dat, jones.dat, wag.dat, mtmam.dat, or your own

        model = 3
                   * models for AAs or codon-translated AAs:
                      * 0:poisson, 1:proportional, 2:Empirical, 3:Empirical+F
                      * 6:FromCodon, 7:AAClasses, 8:REVaa_0, 9:REVaa(nr=189)
        Mgene = 0
                   * AA: 0:rates, 1:separate

        clock = 0  * 0:no clock, 1:global clock; 2:local clock
    fix_alpha = 0  * 0: estimate gamma shape parameter; 1: fix it at alpha
        alpha = 0.5  * initial or fixed alpha, 0:infinity (constant rate)
       Malpha = 0  * different alphas for genes
        ncatG = %s  * # of categories in dG of NSsites models

        getSE = 1  * 0: don't want them, 1: want S.E.s of estimates
 RateAncestor = 1  * (0,1,2): rates (alpha>0) or ancestral states (1 or 2)

   Small_Diff = 1e-6
    cleandata = 0  * remove sites with ambiguity data (1:yes, 0:no)?
*  fix_blength = 0  * 0: ignore, -1: random, 1: initial, 2: fixed
        method = 1   * 0: simultaneous; 1: one branch at a time

''' % (seq_file_name_aa, tree_file_name, ncatG)






for i in os.listdir(infile):
	target_path = infile+'/'+ i +'/'

	to_path = outfile+'/'+i+'/'
	os.mkdir(to_path)
	simulate_file = target_path+'mc.txt'
	with open(simulate_file) as fila:
		seq_dic = {}
		n = 1
		for l in fila:
			k = re.split(r'\s+',l.strip())
			#if k!=[''] and 
			n+=1
			if n >4 and k!=['']:
				#print(k)
				name = k[0]
				seq = ''.join(k[1:])
				seq_dic[name] = seq

	# 写入比对文件

	# fasta
	f_pep = open(to_path+pep_cut_ali,'w')
	for spe in seq_dic:
		f_pep.write('>'+spe+'\n'+seq_dic[spe]+'\n')
	f_pep.close()
	#print(seq_dic.keys())

	# 写入用于跑 paml的aa文件
	f_pep = open(to_path+seq_file_name_aa,'w')
	f_pep.write(str(len(seq_dic))+'  '+str(len(seq_dic[list(seq_dic.keys())[0]]))+'\n\n')
	for spe in seq_dic:
		f_pep.write('>\n'+spe+'\n'+seq_dic[spe]+'\n\n')
	f_pep.close()


	# 写入 aamcl文件
	aa_ctl = open(to_path+aactl_name,'w')
	aa_ctl.write(paml_tmp)
	aa_ctl.close()

	shutil.copy2(dat_file_path,to_path+'my.dat')

	shutil.copy2(target_path+tree_file_name,to_path+tree_file_name)
	# f_pep.close()
			# 写入文件
